#include <algorithm>
#include <atomic>
#include <sstream>
#include <thread>
#include <vector>

#include "tinyndarray.h"

namespace tinyndarray {
size_t get_type_size(Dtype dtype) {
    switch (dtype.tag) {
        case INT8: return sizeof(int8_t);
        case UINT8: return sizeof(uint8_t);
        case INT16: return sizeof(int16_t);
        case UINT16: return sizeof(uint16_t);
        case INT32: return sizeof(int32_t);
        case UINT32: return sizeof(uint32_t);
        case INT64: return sizeof(int64_t);
        case UINT64: return sizeof(uint64_t);
        case FP16: return sizeof(float16_t);
        case FP32: return sizeof(float32_t);
        case FP64: return sizeof(float64_t);
        case QINT8: return sizeof(int8_t);
        case QUINT8: return sizeof(uint8_t);
        case BOOL: return sizeof(bool);
        default: throw std::runtime_error("Unknown data type");
    }
}
// 获取并行参数
void GetParallelParams(int size, int& n_workers, int& n_batch, int& batch_size) {
    n_workers = std::max(1, static_cast<int>(std::thread::hardware_concurrency()));
    batch_size = std::max(1, size / (n_workers * 4));  // 每个批次至少1个元素
    n_batch = (size + batch_size - 1) / batch_size;
}

std::vector<int> ComputeChildSizes(const Shape& shape) {
    const size_t n_shape = shape.size();
    if (n_shape == 0) {
        return {};
    }
    // Compute child sizes from back (the number of children for each dimension)
    std::vector<int> child_sizes(n_shape, 1);
    int size = 1;
    for (size_t depth = n_shape - 1; 0 < depth; depth--) {
        child_sizes[depth] = size;
        size *= shape[depth];
    }
    child_sizes[0] = size;
    return child_sizes;
}

// ----------------------- Utilities for NdArray (Print) -----------------------
static void OutputArrayLine(std::ostream& os, NdArray::ConstIter data, const int size) {
    os << "[";  // Begin of a line
    for (int i = 0; i < size; i++) {
        // 使用解引用操作符 * 获取元素值
        os << *data;

        if (i == size - 1) {
            os << "]";  // End of a line
        } else {
            os << ", ";  // Splitter of an element
        }

        // 移动到下一个元素
        ++data;
    }
}

static void OutputArrayMultiDim(std::ostream& os, NdArray::ConstIter data, const Shape& shape,
                                const std::vector<int>& child_sizes, size_t depth) {
    for (int i = 0; i < shape[depth]; i++) {
        // Heading
        if (i == 0) {
            os << "[";  // begin of array
        } else {
            for (size_t d = 0; d < depth + 1; d++) {  // array indent
                os << " ";
            }
        }

        // Output internal array
        const int& child_size = child_sizes[depth];
        if (depth == shape.size() - 2) {
            // 使用迭代器偏移访问子数组
            OutputArrayLine(os, data, shape[depth + 1]);
        } else {
            OutputArrayMultiDim(os, data, shape, child_sizes, depth + 1);
        }

        // Tailing
        if (i == shape[depth] - 1) {
            os << "]";  // End of array
        } else {
            os << "," << std::endl;  // Splitter of array
        }

        // 移动到下一个子数组的起始位置
        for (int j = 0; j < child_size; j++) {
            ++data;
        }
    }
}

static void OutputNdArray(std::ostream& os, const NdArray& x) {
    const int size = static_cast<int>(x.size());
    const Shape& shape = x.shape();
    const std::vector<int>& child_sizes = ComputeChildSizes(shape);

    if (size == 0 || shape.size() == 0) {
        // Empty
        os << "[]";
    } else if (shape.size() == 1) {
        // 1-dim - 使用 begin() 获取迭代器
        OutputArrayLine(os, x.begin(), size);
    } else {
        // Multi-dim - 使用 begin() 获取迭代器
        OutputArrayMultiDim(os, x.begin(), shape, child_sizes, 0);
    }
}

// 输出运算符
std::ostream& operator<<(std::ostream& os, const NdArray& x) {
    if (x.dtype() != FP16) {
        OutputNdArray(os, x);
    } else {
        OutputNdArray(os, x.astype(FP32));
    }
    return os;
}

template <typename T>
std::string vector_int_to_string(const std::vector<T>& int_vec) {
    std::stringstream ss;
    ss << "[";
    for (size_t i = 0; i < int_vec.size(); i++) {
        ss << int_vec[i];
        if (i < int_vec.size() - 1) {
            ss << ", ";
        }
    }
    ss << "]";
    return ss.str();
}

template std::string vector_int_to_string(const std::vector<int>& int_vec);

std::ostream& operator<<(std::ostream& os, const Shape& shape) {
    os << vector_int_to_string(shape);
    return os;
}

bool is_index_in_range(const Index& index, const Shape& shape) {
    auto nidx = index.size();
    for (auto i = 0; i < nidx; i++) {
        if (index[i] < 0 || index[i] >= shape[i]) return false;
    }
    return true;
}

bool operator==(const Shape& lshape, const Shape& rshape) {
    if (lshape.size() != rshape.size()) {
        return false;
    }
    for (auto idx = 0; idx < lshape.size(); idx++) {
        if (lshape.at(idx) != rshape.at(idx)) {
            return false;
        }
    }
    return true;
}

bool operator!=(const Shape& lshape, const Shape& rshape) {
    return !(lshape == rshape);
}

std::ostream& operator<<(std::ostream& os, const Dtype& dtype) {
    switch (dtype.tag) {
        case INT8: os << "int8";
        case UINT8: os << "uint8";
        case INT16: os << "uint8";
        case UINT16: os << "uint8";
        case INT32: os << "uint8";
        case UINT32: os << "uint8";
        case INT64: os << "uint8";
        case UINT64: os << "uint8";
        case FP16: os << "uint8";
        case FP32: os << "uint8";
        case FP64: os << "uint8";
        case QINT8: os << "uint8";
        case QUINT8: os << "uint8";
        case BOOL: os << "uint8";
        default: throw std::runtime_error("Unknown data type");
    }
    return os;
}

}  // namespace tinyndarray